__doc__ =   """ Module with entities the building blocks of the SMRCA hierachy 
                representation of a macromolecular structure.
            """

from numpy import mean

from Entity import Entity
from MultiEntity import MultiEntity
from ResidueProperties import AA_BACKBONE, AA_SIDECHAIN, AA_NAMES, HOH_NAMES
from AtomProperties import AREAIMOL_VDW_RADII, DEFAULT_AREAIMOL_VDW_RADIUS
from ZenExceptions import ZenException

class Structure(MultiEntity):
    """ The Structure instance contains a collection of Model instances.
    """
    def __init__(self, id, *args, **kwargs):
        self.level = 'S'
        MultiEntity.__init__(self, id, *args, **kwargs)

    def __repr__(self):
        return '<Structure id=%s>' % self.get_id()

    def remove_altmodels(self):
        self._strip_children((0,), 'ne', 'id', forgiving =False)

    def get_dict(self):
        return {'structure':self.get_id()[0]}


class Model(MultiEntity):
    """ The Model instance contains a collection of Chain instances.
    """

    def __init__(self, id, *args, **kwargs):
        self.level = 'M'
        MultiEntity.__init__(self, id, *args, **kwargs)

    def __repr__(self):
        return "<Model id=%s>" % self.get_id()

    def get_dict(self):
        from_parent = self.parent.get_dict()
        from_parent.update({'model':self.get_id()[0]})
        return from_parent


class Chain(MultiEntity):
    def __init__(self, id, *args, **kwargs):
        self.level = 'C'
        MultiEntity.__init__(self, id, *args, **kwargs)

    def remove_hetero(self):
        self._strip_children('H','eq','h_flag', forgiving = False)

    def remove_water(self):
        self._strip_children('H_HOH','eq','name', forgiving = False)

    def residue_count(self):
        return self._count_children('name')

    def residue_freq(self):
        return self._freq_children('name')

    def __repr__(self):
        return "<Chain id=%s>" % self.get_id()

    def get_dict(self):
        from_parent = self.parent.get_dict()
        from_parent.update({'chain_id':self.get_id()[0]})
        return from_parent

class Residue(MultiEntity):
    """ Represents a residue. A Residue object stores atoms.
    """
    def __init__(self, res_long_id, h_flag, seg_id, *args,**kwargs):
        self.level = 'R'
        self.seg_id = seg_id
        self.h_flag = h_flag
        self.res_id = res_long_id[1] #ID number
        self.res_ic = res_long_id[2] #ID long   NAME
        MultiEntity.__init__(self, res_long_id, res_long_id[0], *args, **kwargs)

    def __repr__(self):
        res_name, res_id, res_ic = self.get_id()[0]
        full_name = (res_name, res_id, res_ic)
        return "<Residue %s resseq=%s icode=%s>" % full_name

    #def check_complete(self):
        #""" Checks if backbone and sidechain are complete.
        #"""
        ## consier creating a name -> type mapping in ResidueProperties!
        #if self.name in AA_NAMES:
            #all_sc = set(AA_SIDECHAIN[self.name])
            #all_bb = set(AA_BACKBONE)
        #elif self.name in HOH_NAMES:
            #return (True, True)         # TODO
        #else:
            #return (True, True)         # TODO
        #schain = self._select_children(all_sc, 'contains', 'name', forgiving =False).values()
        #got_sc = set([a.name for a in schain])
        #bbone = self._select_children(all_bb, 'contains', 'name', forgiving =False).values()
        #got_bb = set([a.name for a in bbone])
        #return (got_bb == all_bb, got_sc == all_sc)

    #def remove_sidechain(self, shaveonly =False):
        #schain = list(AA_SIDECHAIN[self.name])  # copy
        #if shaveonly:
            #try:
                #schain.remove(' CB ')   # we leave the Cbeta 'alanine' atom
            #except:
                #pass                    # catch glycins
        #self._strip_children(schain, 'contains', 'name', forgiving =False)

    def remove_hydrogens(self):
        self._strip_children(' H', 'eq', 'element', forgiving =False)

    def get_seg_id(self):
        return self.seg_id

    def set_seg_id(self, seg_id):
        # this does not change the id!
        self.seg_id = seg_id

    def get_res_ic(self):
        return self.res_ic

    def set_res_ic(self, res_ic):
        self.res_ic = res_ic
        self.set_id()

    def get_res_id(self):
        return self.res_id

    def set_res_id(self, res_id):
        self.res_id = res_id
        self.set_id()

    def get_h_flag(self):
        return self.h_flag

    def set_h_flag(self, h_flag):
        if not h_flag in (' ', 'H'):
            raise ZenException, "Only ' ' and 'H' hetero flags allowed."
        if len(self.name) == 3:
            self.name = "%s_%s" % (h_flag, self.name)
        elif len(self.name) == 5:
            self.name = "%s_%s" % (h_flag, self.name[2:])
        else:
            raise ZenException, 'Non-standard residue name'
        self.h_flag = h_flag
        self.set_id()

    def get_id(self):
        return ((self.name, self.res_id, self.res_ic),)

    def get_dict(self):
        from_parent = self.parent.get_dict()
        if self.h_flag != ' ':
            at_type = 'HETATM'
        else:
            at_type = 'ATOM  '
        from_parent.update({'at_type': at_type,
                            'h_flag': self.h_flag,
                            'res_name': self.name,
                            'res_long_id': self.get_id()[0],
                            'res_id': self.res_id,
                            'res_ic': self.res_ic,
                            'seg_id': self.seg_id,})
        return from_parent


class Atom(Entity):
    """ The Atom Entity.
    """
    def __init__(self, at_long_id, at_name, ser_num, coords, occupancy, bfactor, element):
        self.level = 'A'
        self.index = self.hierarchy.index(self.level)
        self.coords = coords
        self.bfactor = bfactor
        self.occupancy = occupancy
        self.ser_num = ser_num
        self.at_id = at_long_id[0]
        self.alt_loc = at_long_id[1]
        self.table = dict([(level,{}) for level in self.hierarchy[self.index+1:]])
        self.element = element
        Entity.__init__(self, at_long_id, at_name)

    def __nonzero__(self):
        return bool(self.id)

    def __repr__(self):
        """ Print Atom object as <Atom atom_name>.
        """
        return "<Atom %s>" % self.get_id()

    def set_name(self, name):
        self.name = name
        self.set_at_id(name.strip())

    def set_at_id(self, at_id):
        self.at_id = at_id
        self.set_id()

    def set_alt_loc(self, alt_loc):
        self.alt_loc = alt_loc
        self.set_id()

    def set_ser_num(self, n):
        self.ser_num = n

    def set_bfactor(self, bfactor):
        self.bfactor = bfactor

    def set_occupancy(self, occupancy):
        self.occupancy = occupancy

    def set_radius(self, radius =None, set =AREAIMOL_VDW_RADII, default =DEFAULT_AREAIMOL_VDW_RADIUS):
        """ Set the atomic radius, defaults to the VDW radius.
        """
        if radius:
            self.radius = radius
        else:
            try:
                self.radius = set[self.element]
            except:
                self.radius = default

    def get_id(self):
        # the id of an atom is not its ' XX ' name
        # but this string after left/right spaces striping.
        return ((self.at_id, self.alt_loc),)

    def get_ser_num(self):
        """ Return atoms serial number
        """
        return self.ser_num

    def get_bfactor(self):
        """ Return B factor.
        """
        return self.bfactor

    def get_occupancy(self):
        """ Return occupancy.
        """
        return self.occupancy

    def get_radius(self):
        """ Return the atomic radius.
        """
        return self.radius

    def get_dict(self):
        from_parent = self.parent.get_dict()
        from_parent.update({'at_name': self.name,
                            'ser_num': self.ser_num,
                            'coords': self.coords,
                            'occupancy': self.occupancy,
                            'bfactor': self.bfactor,
                            'alt_loc': self.alt_loc,
                            'at_long_id': self.get_id()[0],
                            'at_id': self.at_id,
                            'element': self.element})
        return from_parent

#EOF
